"""extract_features.py: 从训练好的β-VAE模型中提取图像的解耦特征"""

import os
import argparse
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from tqdm import tqdm

from model import BetaVAE_H, BetaVAE_B
from dataset import CustomImageFolder, CustomTensorDataset
from utils import cuda


def load_model(args):
    """加载训练好的β-VAE模型"""
    # 根据参数选择模型架构
    if args.model == 'H':
        model = BetaVAE_H(z_dim=args.z_dim, nc=args.nc)
    elif args.model == 'B':
        model = BetaVAE_B(z_dim=args.z_dim, nc=args.nc)
    else:
        raise ValueError("模型类型仅支持 'H' 或 'B'")

    # 加载权重
    if not os.path.exists(args.ckpt_path):
        raise FileNotFoundError(f"权重文件不存在: {args.ckpt_path}")
    checkpoint = torch.load(args.ckpt_path, map_location=torch.device('cpu' if not args.cuda else 'cuda'))
    model.load_state_dict(checkpoint['model_states'])

    # 切换到评估模式
    model.eval()
    return cuda(model, args.cuda)


def get_data_loader(args):
    """创建数据加载器"""
    # 图像预处理（与训练时保持一致）
    transform = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
    ])

    # 加载数据集
    if args.dataset.lower() == 'dsprites':
        # dSprites是预处理好的npz文件
        data = np.load(args.data_dir, encoding='bytes')
        data = torch.from_numpy(data['imgs']).unsqueeze(1).float()
        dataset = CustomTensorDataset(data_tensor=data)
    else:
        # CelebA/3DChairs是图像文件夹
        dataset = CustomImageFolder(root=args.data_dir, transform=transform)

    # 创建数据加载器
    return DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True
    )


def extract_features(args, model, data_loader):
    """提取解耦特征（mu向量）"""
    all_features = []
    with torch.no_grad():  # 禁用梯度计算
        for x in tqdm(data_loader, desc="提取特征"):
            x = cuda(x, args.cuda)
            # 编码得到分布参数
            distributions = model._encode(x)
            mu = distributions[:, :args.z_dim]  # 提取均值作为解耦特征
            all_features.append(mu.cpu().numpy())

    # 拼接所有特征
    return np.concatenate(all_features, axis=0)


def main(args):
    # 设置随机种子（保证可复现性）
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    # 自动推断图像通道数（1 for dSprites, 3 for others）
    args.nc = 1 if args.dataset.lower() == 'dsprites' else 3

    # 加载模型
    print(f"加载模型权重: {args.ckpt_path}")
    model = load_model(args)

    # 加载数据
    print(f"加载数据集: {args.data_dir}")
    data_loader = get_data_loader(args)

    # 提取特征
    print("开始提取解耦特征...")
    features = extract_features(args, model, data_loader)
    print(f"特征提取完成，形状: {features.shape} (样本数 × 潜在维度)")

    # 保存特征
    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
    np.save(args.save_path, features)
    print(f"特征已保存至: {args.save_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="从β-VAE模型提取解耦特征")

    # 模型参数（需与训练时保持一致）
    parser.add_argument('--model', default='H', type=str, choices=['H', 'B'],
                        help="模型类型 (H: Higgins et al. / B: Burgess et al.)")
    parser.add_argument('--z_dim', default=10, type=int,
                        help="潜在空间维度 (与训练时一致)")
    parser.add_argument('--ckpt_path', default=r'D:\Code\deep-learning-code\Beta-VAE\checkpoints\main\last', type=str,
                        help="模型权重文件路径 (如 checkpoints/celeba_H_beta10_z10/last.pth)")

    # 数据参数
    parser.add_argument('--dataset', default='CelebA', type=str,
                        help="数据集名称")
    parser.add_argument('--data_dir', default=r'D:\Code\2-ZSL\Zero-Shot-Learning\data\0HP', type=str,
                        help="数据集路径 (dSprites为npz文件路径，其他为图像文件夹路径)")
    parser.add_argument('--batch_size', default=64, type=int, help="批处理大小")
    parser.add_argument('--num_workers', default=0, type=int, help="数据加载线程数")

    # 其他参数
    parser.add_argument('--save_path', default='outputs', type=str,
                        help="特征保存路径 (如 outputs/decoupled_features.npy)")
    parser.add_argument('--cuda', default=True,
                        help="是否使用GPU")
    parser.add_argument('--seed', default=1, type=int, help="随机种子")

    args = parser.parse_args()
    main(args)